import mmengine
import pickle
import json
from .base_det_dataset import BaseDetDataset
from mmdet.registry import DATASETS


@DATASETS.register_module()
class MyDataset(BaseDetDataset):
    METAINFO = {
        'classes': ('tampered'),
        'palette': [(220, 20, 60)]
    }

    def load_data_list(self, ):
        with open(self.ann_file, 'rb') as f:
            ann_list = pickle.load(f)

        data_infos = []
        for i, (k, v) in enumerate(ann_list.items()):

            width = v['w']
            height = v['h']
            instances = []

            for anns in v['b']:

                # print('anns' + str(anns))
                # cr
                # 假设 anns 有多余信息，只取前四个作为 bbox
                if len(anns) < 4:
                    print(f"Warning: invalid bbox format for image {k}, skipping annotation.")
                    continue

                # 提取 bbox 的前四个坐标值
                x_min, y_min, x_max, y_max = anns[:4]
                if (x_min < 0 or y_min < 0 or x_max > width or y_max > height or
                        x_max <= x_min or y_max <= y_min):
                    print(f"Warning: skipping invalid bbox for image {k}: {x_min, y_min, x_max, y_max}")
                    continue  # 跳过不合法标注

                instance = {}
                instance['bbox'] = list(anns)
                instance['bbox_label'] = 0
                instance['bbox_whs'] = (width, height)
                instance['bbox_pts'] = list(anns)
                instances.append(instance)

            # 如果有合法的标注，才将图像信息添加到 data_infos 中
            if instances:
                data_infos.append(
                    dict(
                        img_path=k,
                        img_id=i,
                        width=width,
                        height=height,
                        instances=instances
                    ))

        return data_infos
